package org.liang.code.utils.pdf;

import com.lowagie.text.DocumentException;
import com.lowagie.text.pdf.*;

import java.awt.*;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * 水印工具
 */
public class WaterMakerUtil {

    private final PdfReader reader;
    private final ByteArrayOutputStream outputStream;
    private final PdfStamper stamp;
    private final String text;
    private final int fontsize;
    private final float opacity;
    private final int textDensity;
    private final BaseFont font = PDFUtils.CJKFont_H;
    private Color color = Color.BLACK;
    private float rotation = 45F;

    /**
     * Watermarker的构造函数
     * @param input PDF的输入流
     * @param text 水印文字
     * @param fontsize 文字大小
     * @param opacity 文字透明度
     */
    public WaterMakerUtil(byte[] input, String text, int fontsize, float opacity, int textDensity) throws IOException, DocumentException {
        this.reader = new PdfReader(input);
        this.outputStream = new ByteArrayOutputStream();
        // 将额外内容应用于 PDF 文档的页面
        this.stamp = new PdfStamper(reader, outputStream);
        this.text = text;
        this.fontsize = fontsize;
        this.opacity = opacity;
        this.textDensity = textDensity;
    }

    /**
     * 水印文字颜色
     * @param color 默认为黑色
     * @return Watermarker
     */
    public WaterMakerUtil withColor(Color color) {
        this.color = color;
        return this;
    }


    /**
     * 倾斜角度
     * @param rotation 默认为45度
     * @return Watermarker
     */
    public WaterMakerUtil withRotation(float rotation) {
        this.rotation = rotation;
        return this;
    }


    public byte[] write() throws IOException, DocumentException {
        final BaseFont bf = (font != null) ? font : PDFUtils.CJKFont_H;
        // 获取文档所有页数
        int pagecount = reader.getNumberOfPages();
        PdfGState gstate  = new PdfGState();
        // 设置当前填充 Alpha 常数，指定要用于透明成像模型中填充作的常量形状或恒定不透明度值
        gstate.setFillOpacity(opacity); // 设置填充透明度
        gstate.setStrokeOpacity(opacity); // 设置描边透明度

        for (int i = 1; i <= pagecount; i++) {
            // 使用OverContent以确保水印在页面内容之上
            PdfContentByte content = stamp.getOverContent(i);

            // 设置水印透明度
            content.saveState();
            content.setGState(gstate);
            // 设置字体和颜色
            content.setFontAndSize(bf, fontsize);
            content.setColorFill(color);


            // 获取页宽
            float width = content.getPdfDocument().getPageSize().getWidth();
            // 获取页高
            float height = content.getPdfDocument().getPageSize().getHeight();
            // 计算文字的密度
            int xStep = (int) (width / Math.sqrt(textDensity));
            int yStep = (int) (height / Math.sqrt(textDensity));

            // 在页面上绘制文字
            for (int x = fontsize; x < width; x += xStep) {
                for (int y = fontsize; y < height; y += yStep) {
                    content.beginText();
                    content.showTextAligned(PdfContentByte.ALIGN_CENTER, text, x, y, rotation); // 45度倾斜
                    content.endText();
                }
            }
            // 恢复状态
            content.restoreState();
        }
        stamp.close();
        return outputStream.toByteArray();
    }
}
